import os
import sys
import time
sys.path.append("/usr/local/lib/python3.8/site-packages")
sys.path.append("/opt/openrobots/lib/python3.8/site-packages") # 防止找不到pinocchio和crocoddyl，这个需要更改成实际的地址
##crocoddyl pinocchio安装方法--
# https://mp.weixin.qq.com/s/G6sgz2IIpc359moaoBq3iQ
from pinocchio import visualize
import numpy as np
import pinocchio
from pinocchio import RobotWrapper
import crocoddyl
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation as R

## feidedaoRobot 飞的岛

current_directory = os.getcwd()
parent_directory = os.path.dirname(current_directory)
print("路径：", current_directory)
# 定义模型路径和文件名

modelPath = current_directory+'/'
URDF_FILENAME = "robot_description/urdf/robot_rviz.urdf"

# 加载 URDF 文件
rrobot = RobotWrapper.BuildFromURDF(modelPath + URDF_FILENAME, [modelPath], pinocchio.JointModelFreeFlyer())
rmodel = rrobot.model
rdata = rmodel.createData()

# 定义右脚和左脚的 ID
rightFoot = 'right_ankle_roll_link'
leftFoot = 'left_ankle_roll_link' #  "base_link"
display = crocoddyl.MeshcatDisplay(rrobot, frameNames=[rightFoot, leftFoot])

dd1 = np.genfromtxt('robot_state.txt', delimiter=' ', skip_footer=0)[:, :]
qpos = dd1[:,:rmodel.nq].copy()
qvel = dd1[:,rmodel.nq:rmodel.nq+rmodel.nv].copy()
baseQuat = qpos[:, 3:7].copy()
baseQuat[:, 0] = qpos[:, 4].copy() # x
baseQuat[:, 1] = qpos[:, 5].copy() # y
baseQuat[:, 2] = qpos[:, 6].copy() # z
baseQuat[:, 3] = qpos[:, 3].copy() # w

footPosMujoco = dd1[:, rmodel.nq+rmodel.nv:rmodel.nq+rmodel.nv+3].copy()
footQuat = dd1[:, rmodel.nq+rmodel.nv+3:rmodel.nq+rmodel.nv+7].copy()
footQuat[:, 0] = dd1[:, rmodel.nq+rmodel.nv+4].copy() # x
footQuat[:, 1] = dd1[:, rmodel.nq+rmodel.nv+5].copy() # y
footQuat[:, 2] = dd1[:, rmodel.nq+rmodel.nv+6].copy() # z
footQuat[:, 3] = dd1[:, rmodel.nq+rmodel.nv+3].copy() # w
r1 = R.from_quat(footQuat)  #xyzw
FootEulerMujoco = r1.as_euler('ZYX')


footVelMujoco = dd1[:, rmodel.nq+rmodel.nv+7:rmodel.nq+rmodel.nv+13].copy() # w v
num = len(qpos)


# 微分计算角速度
footW_diff = np.zeros((num,3))
dt = 0.002
for i in range(num-1):
    r0 = R.from_quat(footQuat[i, :].copy())
    r1 = R.from_quat(footQuat[i+1, :].copy())
    dq = r1 * r0.inv()
    footW_diff[i+1, :] = (dq.as_rotvec() / dt)  # 角速度是旋转向量除以时间间隔
# 微分计算速度
footV_diff = np.zeros((num, 3))
footV_diff[1:, :] = (footPosMujoco[1:, :].copy() - footPosMujoco[0:-1, :].copy())/dt

# 将mujoco的速度世界坐标系，转换成base坐标系
rbase = R.from_quat(baseQuat) # xyzw
Rbase = rbase.as_matrix()
baseLocalV = qvel[:, :3].copy()
for i in range(num):
    baseLocalV[i, :] = Rbase[i, :, :].transpose() @ qvel[i, :3].copy()

if 0: # 需要观看动画就把这个打开 1
    start = 1
    for j in range(num - start):
        i = j+start
        q0 = qpos[i, :].copy()
        q0[3] = qpos[i, 4].copy()
        q0[4] = qpos[i, 5].copy()
        q0[5] = qpos[i, 6].copy()
        q0[6] = qpos[i, 3].copy()
        display.display([q0])
        time.sleep(0.002)
        print(i)

# 计算脚的位置、姿态和角速度
posFoot = []
angleEulerFoot = []
angularVelFoot = []
linearVelFoot = []
for i in range(num):
    qpin = qpos[i, :].copy()
    qpin[3] = qpos[i, 4].copy()
    qpin[4] = qpos[i, 5].copy()
    qpin[5] = qpos[i, 6].copy()
    qpin[6] = qpos[i, 3].copy()

    vpin = qvel[i, :].copy()
    vpin[:3] = baseLocalV[i, :].copy() # local velocity
    # vpin[3:6] = baseLocalW[i, :].copy()

    rdata = rmodel.createData()
    pinocchio.forwardKinematics(rmodel, rdata, qpin, vpin)
    pinocchio.updateFramePlacements(rmodel, rdata)

    foot_frame_id = rmodel.getFrameId(leftFoot)
    foot_pose = rdata.oMf[foot_frame_id]
    foot_position = foot_pose.translation
    foot_rotation = foot_pose.rotation

    # 注意这里ZYX是大写，yaw，pitch，roll
    r = R.from_matrix(foot_rotation)
    euler_angle = r.as_euler('ZYX')

    # 获取速度
    v = pinocchio.getFrameVelocity(rmodel, rdata, foot_frame_id, pinocchio.LOCAL_WORLD_ALIGNED)
    linear_velocity = v.linear
    angular_velocity = v.angular

    # 通过雅可比的方法，计算速度
    # Jf = pinocchio.computeFrameJacobian(rmodel, rdata, qpin, foot_frame_id, pinocchio.LOCAL_WORLD_ALIGNED)
    # v = Jf @ vpin
    # linear_velocity = v[:3]
    # angular_velocity = v[3:]

    posFoot.append(foot_position)
    angleEulerFoot.append(euler_angle)
    angularVelFoot.append(angular_velocity)
    linearVelFoot.append(linear_velocity)

posFoot = np.array(posFoot)
angleEulerFoot = np.array(angleEulerFoot)
angularVelFoot = np.array(angularVelFoot)
linearVelFoot = np.array(linearVelFoot)

for i in range(3):
    plt.figure(i + 10)
    plt.plot(angleEulerFoot[:, i], 'y', label="euler py")
    plt.plot(FootEulerMujoco[:, i], 'r', label="euler mujoco")
    plt.title(i)
    plt.legend()
    plt.show()

plt.close('all')
for i in range(3):
    plt.figure(i + 110)
    plt.plot(linearVelFoot[:, i], 'k', label="foot v py")
    plt.plot(footVelMujoco[:, i+3], 'r', label="foot v mujoco")
    plt.plot(footV_diff[:, i], 'y', label="foot v diff")
    # plt.plot(baseLocalV[:, i], 'm', label="foot m q")
    # plt.plot(qvel[:, i], 'b', label="foot m q")
    plt.title(i)
    plt.legend()
    plt.show()


plt.close('all')
for i in range(3):
    plt.figure(i + 20)
    plt.plot(angularVelFoot[:, i], 'k', label="foot w py")
    plt.plot(footVelMujoco[:, i], 'r', label="foot w mujoco")
    plt.plot(footW_diff[:, i], 'y', label="foot w diff")
    # plt.plot(qvel[:, i+3], 'b', label="foot qvel q")
    plt.title(i)
    plt.legend()
    plt.show()

plt.close('all')
for i in range(3):
    plt.figure(i + 30)
    plt.plot(posFoot[:, i], 'k', label="foot p py")
    plt.plot(footPosMujoco[:, i], 'r', label="foot p mujoco")
    plt.title(i)
    plt.legend()
    plt.show()

input("Press Enter to continue...")